package com.shortlink.common.ratelimiter;

import com.shortlink.common.exception.ServiceException;
import lombok.*;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.util.Collections;

import static com.shortlink.common.constant.RedisKeyConstant.COUNT_RATE_LIMITER_KEY;

/**
 * 计数器法限流器实现（固定窗口法）
 *
 * @author wp
 */
public record CounterRateLimiter(StringRedisTemplate stringRedisTemplate,
                                 Config config) implements RateLimiter {
    /**
     * 限流器lua脚本
     */
    private static final String RATE_LIMITER_LUA_SCRIPT_STR =
            """
                    local accessKey = KEYS[1]
                    local timeWindow = tonumber(ARGV[1])

                    -- 原子递增访问次数，并获取递增后的值
                    local currentAccessCount = redis.call("INCR", accessKey)

                    -- 设置键的过期时间
                    if currentAccessCount == 1 then
                        redis.call("EXPIRE", accessKey, timeWindow)
                    end

                    -- 返回当前访问次数
                    return currentAccessCount""";
    //private static final String RATE_LIMITER_SCRIPT_PATH = "/count_rate_limiter.lua";
    /**
     * 限流器lua脚本
     */
    private static final RedisScript<Long> RATE_LIMITER_SCRIPT;

    static {
        //RATE_LIMITER_SCRIPT = new DefaultRedisScript<>(RATE_LIMITER_SCRIPT_PATH, Long.class);
        RATE_LIMITER_SCRIPT = new DefaultRedisScript<>(RATE_LIMITER_LUA_SCRIPT_STR, Long.class);
    }

    @Override
    public boolean tryAcquire(String name) {
        String redisKey = String.format(COUNT_RATE_LIMITER_KEY, name);
        Long count = stringRedisTemplate.execute(RATE_LIMITER_SCRIPT,
                Collections.singletonList(redisKey), Collections.EMPTY_LIST, config.getWindowSize());
        if (count == null) {
            throw new ServiceException("lua脚本执行异常");
        }
        return count <= config.getLimit();
    }

    @Data
    public static class Config {
        /**
         * 窗口时长，毫秒级
         */
        private Long windowSize;
        /**
         * 限流大小
         */
        private Integer limit;
    }
}